<template>
    <div class="generator-view">
        <div class="content">
            <!-- 左侧参数面板 -->
            <div class="params-panel">
                <div class="model-select">
                    <span class="model-name">Flux</span>
                    <el-dropdown>
                        <span class="dropdown-link">
                            <el-icon><ArrowDown /></el-icon>
                        </span>
                        <template #dropdown>
                            <el-dropdown-menu>
                                <el-dropdown-item>SDXL</el-dropdown-item>
                                <el-dropdown-item>ControlNet</el-dropdown-item>
                            </el-dropdown-menu>
                        </template>
                    </el-dropdown>
                </div>

                <!-- 提示词输入 -->
                <div class="prompt-input">
                    <el-input
                        v-model="params.prompt"
                        type="textarea"
                        :rows="6"
                        :placeholder="'在此输入提示词...'"
                        resize="none"
                    />
                </div>

                <!-- 图片尺寸设置 -->
                <div class="size-settings">
                    <div class="setting-item">
                        <span class="label">宽度</span>
                        <el-input-number
                            v-model="params.width"
                            :min="512"
                            :max="2048"
                            :step="64"
                            controls-position="right"
                        />
                    </div>
                    <div class="setting-item">
                        <span class="label">高度</span>
                        <el-input-number
                            v-model="params.height"
                            :min="512"
                            :max="2048"
                            :step="64"
                            controls-position="right"
                        />
                    </div>
                </div>

                <!-- 随机种子 -->
                <div class="seed-input">
                    <div class="setting-header">
                        <span class="label">随机种子</span>
                        <el-button 
                            type="primary" 
                            link 
                            @click="params.seed = Math.floor(Math.random() * 1000000000)"
                        >
                            随机
                        </el-button>
                    </div>
                    <el-input
                        v-model="params.seed"
                        placeholder="输入种子值"
                    >
                        <template #append>
                            <el-button @click="copyToClipboard(params.seed.toString())">
                                <el-icon><CopyDocument /></el-icon>
                            </el-button>
                        </template>
                    </el-input>
                </div>

                <!-- 生成按钮 -->
                <div class="generate-button">
                    <el-button type="primary" :loading="generating" @click="handleGenerate">
                        {{ generating ? '生成中...' : '生成 (耗时约20秒)' }}
                    </el-button>
                </div>
            </div>

            <!-- 右侧预览区域 -->
            <div class="preview-panel">
                <div class="preview-container" :class="{ generating }">
                    <template v-if="generating">
                        <el-progress type="circle" :percentage="progress" />
                        <p class="progress-text">正在生成中...</p>
                    </template>
                    <template v-else-if="previewUrl">
                        <el-image :src="previewUrl" fit="contain" />
                        <div class="preview-actions">
                            <el-button-group>
                                <el-button type="primary">
                                    <el-icon><Download /></el-icon>
                                    下载
                                </el-button>
                                <el-button type="primary">
                                    <el-icon><Share /></el-icon>
                                    分享
                                </el-button>
                            </el-button-group>
                        </div>
                    </template>
                    <template v-else>
                        <el-empty description="点击生成按钮开始创作" />
                    </template>
                </div>
            </div>
        </div>
    </div>
</template>

<script setup lang="ts">
import { ref, onMounted, onBeforeUnmount } from 'vue'
import { useRoute } from 'vue-router'
import { ArrowDown, CopyDocument, Download, Share } from '@element-plus/icons-vue'
import { ElMessage } from 'element-plus'
import axios from 'axios'

const route = useRoute()
const ws = ref<WebSocket | null>(null)

const params = ref({
    prompt: '',
    width: 1024,
    height: 1024,
    seed: Math.floor(Math.random() * 1000000000)
})

const generating = ref(false)
const progress = ref(0)
const previewUrl = ref('')
const promptId = ref('')

const copyToClipboard = async (text: string) => {
    try {
        await navigator.clipboard.writeText(text)
        ElMessage.success('已复制到剪贴板')
    } catch (err) {
        ElMessage.error('复制失败')
    }
}

// 连接 WebSocket
const connectWebSocket = () => {
    ws.value = new WebSocket(import.meta.env.VITE_COMFYUI_WS_URL)
    
    ws.value.onopen = () => {
        console.log('WebSocket connected')
    }

    ws.value.onerror = (error) => {
        console.error('WebSocket error:', error)
        ElMessage.error('连接失败，请检查服务器状态')
    }

    ws.value.onmessage = (event) => {
        const data = JSON.parse(event.data)
        
        if (data.type === 'progress') {
            progress.value = Math.round(data.data.value * 100)
        } else if (data.type === 'executed') {
            if (data.data.node === '10' && data.data.output) {
                const imageUrl = `${import.meta.env.VITE_COMFYUI_API_URL}/view?filename=${data.data.output.images[0].filename}&subfolder=&type=temp`
                previewUrl.value = imageUrl
                generating.value = false
            }
        }
    }
}

const maxRetries = 3
const retryDelay = 2000

const connectWebSocketWithRetry = async (retries = 0) => {
    try {
        connectWebSocket()
    } catch (error) {
        if (retries < maxRetries) {
            console.log(`重试连接 WebSocket (${retries + 1}/${maxRetries})`)
            setTimeout(() => {
                connectWebSocketWithRetry(retries + 1)
            }, retryDelay)
        } else {
            ElMessage.error('无法连接到服务器，请刷新页面重试')
        }
    }
}

onMounted(() => {
    connectWebSocketWithRetry()
    
    if (route.query.template) {
        // 从模板加载参数
        const { prompt, width, height } = route.query
        params.value = {
            prompt: prompt as string || '',
            width: Number(width) || 1024,
            height: Number(height) || 1024,
            seed: Math.floor(Math.random() * 1000000000)
        }
    }
})

onBeforeUnmount(() => {
    if (ws.value) {
        ws.value.close()
    }
})

const handleGenerate = async () => {
    generating.value = true
    progress.value = 0
    previewUrl.value = ''

    try {
        // 构建工作流
        const workflow = {
            "3": {
                "inputs": {
                    "seed": params.value.seed,
                    "steps": 20,
                    "cfg": 7,
                    "sampler_name": "euler_a",
                    "scheduler": "normal",
                    "denoise": 1,
                    "model": ["4", 0],
                    "positive": ["6", 0],
                    "negative": ["7", 0],
                    "latent_image": ["5", 0]
                },
                "class_type": "KSampler"
            },
            "4": {
                "inputs": {
                    "ckpt_name": "flux_v1.safetensors"
                },
                "class_type": "CheckpointLoaderSimple"
            },
            "5": {
                "inputs": {
                    "width": params.value.width,
                    "height": params.value.height,
                    "batch_size": 1
                },
                "class_type": "EmptyLatentImage"
            },
            "6": {
                "inputs": {
                    "text": params.value.prompt
                },
                "class_type": "CLIPTextEncode"
            },
            "7": {
                "inputs": {
                    "text": "low quality, bad anatomy, worst quality"
                },
                "class_type": "CLIPTextEncode"
            },
            "8": {
                "inputs": {
                    "samples": ["3", 0],
                    "vae": ["4", 2]
                },
                "class_type": "VAEDecode"
            },
            "10": {
                "inputs": {
                    "images": ["8", 0]
                },
                "class_type": "PreviewImage"
            },
            "11": {
                "inputs": {
                    "filename_prefix": "ComfyUI",
                    "images": ["8", 0]
                },
                "class_type": "SaveImage"
            }
        }

        // 发送生成请求
        const response = await axios.post(import.meta.env.VITE_COMFYUI_API_URL + '/prompt', {
            prompt: workflow
        })

        promptId.value = response.data.prompt_id

        // 等待生成完成
        const checkInterval = setInterval(() => {
            if (!generating.value) {
                clearInterval(checkInterval)
            }
        }, 1000)
    } catch (error) {
        console.error('生成失败:', error)
        ElMessage.error('生成失败，请重试')
        generating.value = false
        progress.value = 0
    }
}
</script>

<style scoped>
.generator-view {
    height: 100%;
    background-color: #fff;
    overflow: hidden;
}

.content {
    display: flex;
    height: 100%;
    overflow: hidden;
}

.params-panel {
    width: 400px;
    padding: 20px;
    border-right: 1px solid #e6e6e6;
    display: flex;
    flex-direction: column;
    gap: 20px;
    overflow-y: auto;
}

.model-select {
    display: flex;
    align-items: center;
    justify-content: space-between;
    padding: 12px;
    background-color: #f5f7fa;
    border-radius: 8px;
}

.model-name {
    font-size: 16px;
    font-weight: bold;
}

.dropdown-link {
    cursor: pointer;
    color: #409eff;
}

.prompt-input {
    flex-grow: 0;
}

.size-settings {
    display: flex;
    gap: 12px;
}

.setting-item {
    flex: 1;
    display: flex;
    flex-direction: column;
    gap: 8px;
}

.setting-header {
    display: flex;
    justify-content: space-between;
    align-items: center;
}

.label {
    font-size: 14px;
    color: #606266;
}

.generate-button {
    margin-top: auto;
}

.preview-panel {
    flex: 1;
    padding: 20px;
    display: flex;
    flex-direction: column;
    overflow: hidden;
}

.preview-container {
    flex: 1;
    display: flex;
    flex-direction: column;
    justify-content: center;
    align-items: center;
    background-color: #f5f7fa;
    border-radius: 8px;
    position: relative;
    overflow: hidden;
}

.preview-container.generating {
    background-color: rgba(0, 0, 0, 0.02);
}

.progress-text {
    margin-top: 12px;
    color: #909399;
}

.preview-actions {
    position: absolute;
    bottom: 20px;
    left: 50%;
    transform: translateX(-50%);
}

:deep(.el-input-number) {
    width: 100%;
}

:deep(.el-progress) {
    width: 120px;
    height: 120px;
}

:deep(.el-image) {
    max-width: 100%;
    max-height: 100%;
}
</style> 